from typing import Dict, Any, List

import numpy as np
import pandas as pd


def convert_ID2num(originID_l: List[str]) -> Dict[str, int]:
    ID2num_m = {}
    count = 1
    for ID in originID_l:
        if ID not in ID2num_m.keys():
            ID2num_m[ID] = count
            count += 1
    return ID2num_m


# assume it has three colum userID, itemID, rating
def convert_ratings(rating: pd.DataFrame) -> pd.DataFrame:
    userID_l = rating['userID'].tolist()
    user2num_m = convert_ID2num(userID_l)

    itemID_l = rating['itemID'].tolist()
    item2num_m = convert_ID2num(itemID_l)

    rating['itemID'] = rating['itemID'].apply(lambda x: item2num_m[x])
    rating['userID'] = rating['userID'].apply(lambda x: user2num_m[x])

    return rating


def stats_rating(rating_after_convert: pd.DataFrame):
    n_user = np.max(rating_after_convert["userID"])
    n_item = np.max(rating_after_convert["itemID"])
    sizeof_float = 4
    print(
        f'n_user {n_user}, n_item {n_item}, index size {n_user * n_item * sizeof_float / 1024 / 1024 / 1024 / 1024:.2f}TB')
